import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from model import LSTMModel

# 加载数据并准备测试集
def load_test_data(csv_file):
    df = pd.read_csv(csv_file)

    # 提取输入特征 (A1, A2, ..., B120) 和标签 (match_label)
    X = df.drop(columns=["match_label"]).values
    y = df["match_label"].apply(lambda x: 1 if x == "team1" else 0).values

    # 标准化特征
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    # 拆分训练集和测试集 (这里只用测试集)
    _, X_test, _, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 将数据转为 PyTorch tensor
    X_test = torch.tensor(X_test, dtype=torch.float32).reshape(-1, 1, 360)
    y_test = torch.tensor(y_test, dtype=torch.long)

    test_dataset = torch.utils.data.TensorDataset(X_test, y_test)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)

    return test_loader

# 模型评估函数
def evaluate_model(model, test_loader, criterion):
    model.eval()  # 设置模型为评估模式
    correct = 0
    total = 0
    total_loss = 0.0

    with torch.no_grad():
        for data, labels in test_loader:
            data, labels = data.to(device), labels.to(device)

            # 前向传播
            outputs = model(data)
            loss = criterion(outputs, labels)
            total_loss += loss.item()

            # 预测类别
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = correct / total
    avg_loss = total_loss / len(test_loader)

    return accuracy, avg_loss

if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载训练好的模型
    model = LSTMModel(input_size=360, hidden_size=128, num_layers=2, output_size=2).to(device)
    model.load_state_dict(torch.load("lstm_model.pth"))

    # 加载测试数据
    csv_file = 'matches_odds.csv'
    test_loader = load_test_data(csv_file)

    # 初始化损失函数
    criterion = nn.CrossEntropyLoss()

    # 评估模型
    accuracy, avg_loss = evaluate_model(model, test_loader, criterion)

    print(f"测试集准确率: {accuracy * 100:.2f}%")
    print(f"测试集平均损失: {avg_loss:.4f}")


